package loveqq.jdbc.log.replace;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.SQLReplaceable;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.*;
import com.alibaba.druid.sql.visitor.SQLASTVisitor;

import java.math.BigDecimal;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.time.format.DateTimeFormatter;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.SimpleTimeZone;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 使用Druid实现Sql解析及替换
 */
public class DruidSqlReplace implements SqlReplace {

    @Override
    public String replace(String sql, Map<Integer, Object> paramMap, String dbType) throws Exception {
        // 解析SQL生成AST根节点
        List<SQLStatement> statements = SQLUtils.parseStatements(sql, dbType);
        // 参数位置
        AtomicInteger paramIndex = new AtomicInteger(0);
        // AST节点遍历器：进行参数值替换
        SQLASTVisitor visitor = new GenericASTVisitorAdapter() {
            @Override
            public boolean visit(SQLVariantRefExpr expr) {
                if (expr != null && "?".equals(expr.toString())) {
                    // 参数位置加1
                    int index = paramIndex.incrementAndGet();
                    // 如果AST父节点是可替换，就替换为参数
                    SQLObject parent = expr.getParent();
                    if (parent instanceof SQLReplaceable) {
                        Object param = paramMap.get(index);
                        // 替换占位符为参数值
                        ((SQLReplaceable) parent).replace(expr, wrapperParam(param, parent));
                    }
                }
                return false;
            }
        };

        // 执行SQL替换
        return statements.stream()
                .map(statement -> {
                    statement.accept(visitor);
                    // 返回替换了参数值的SQL
                    return statement.toString();
                })
                .collect(Collectors.joining("\n\n"));
    }

    /**
     * 包装参数，防止SQL注入
     */
    private SQLExpr wrapperParam(Object param, SQLObject parent) {
        SQLExpr sqlExpr;

        if (param == null) {
            sqlExpr = new SQLNullExpr();
        } else if (param instanceof BigDecimal) {
            sqlExpr = new SQLDecimalExpr((BigDecimal) param);
        } else if (param instanceof Number) {
            sqlExpr = new SQLNumberExpr((Number) param);
        } else if (param instanceof String) {
            sqlExpr = new SQLCharExpr((String) param);
        } else if (param instanceof Boolean) {
            sqlExpr = new SQLBooleanExpr((Boolean) param);
        } else if (param instanceof java.sql.Date) {
            Date date = new Date(((java.sql.Date) param).getTime());
            sqlExpr = new SQLDateExpr(date, SimpleTimeZone.getDefault());
        } else if (param instanceof java.sql.Time) {
            Date time = new Date(((java.sql.Time) param).getTime());
            sqlExpr = new SQLTimeExpr(time, SimpleTimeZone.getDefault());
        } else if (param instanceof java.sql.Timestamp) {
            Date timestamp = new Date(((java.sql.Timestamp) param).getTime());
            sqlExpr = new SQLTimestampExpr(timestamp, SimpleTimeZone.getDefault());
        } else if (param instanceof Date) {
            sqlExpr = new SQLTimestampExpr((Date) param, SimpleTimeZone.getDefault());
        } else if (param instanceof LocalDate) {
            sqlExpr = new SQLDateExpr(((LocalDate) param).format(DateTimeFormatter.ofPattern("yyyy-MM-dd")));
        } else if (param instanceof LocalTime) {
            sqlExpr = new SQLTimeExpr(((LocalTime) param).format(DateTimeFormatter.ofPattern("HH:mm:ss")));
        } else if (param instanceof LocalDateTime) {
            sqlExpr = new SQLTimestampExpr(((LocalDateTime) param).format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        } else {
            sqlExpr = new SQLCharExpr("\n************ 暂不支持格式, 需要手动设置 ************\n");
        }
        // 设置AST父节点，返回
        sqlExpr.setParent(parent);
        return sqlExpr;
    }
}
